#include "elfparser.h"
#include "utils.h"
#include <assert.h>
#include <common.h>
#include <ftrace.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

#define SPACE_BUF_MAX_LEN 512
static char space_buffer[SPACE_BUF_MAX_LEN];
static char* spaces(int n) {
    assert(n >= 0 && n < SPACE_BUF_MAX_LEN);
    memset(space_buffer, ' ', n);
    space_buffer[n] = '\0';
    return space_buffer;
}

#define MAX_FUNC_ITEMS  1024
#define MAX_FTRACE_ITEMS  1000

static FuncEntry func_table[MAX_FUNC_ITEMS];
static int func_count = 0;

static FrameEntry frame_table[MAX_FTRACE_ITEMS];

static int front = 0, rear = 0;
static size_t depth = 0;

static FuncEntry* find_func(vaddr_t target) {
    for (int i = 0; i < func_count; ++i) {
        if (func_table[i].start <= target && target < func_table[i].start + func_table[i].size)
            return &func_table[i];
    }
    return NULL;
}

/* 添加 FrameEntry */
static void append_frame(vaddr_t src_addr, vaddr_t dst_addr, enum FrameType type, size_t depth) {
    assert(src_addr != 0 && dst_addr != 0);
    if ( (rear+1) % MAX_FTRACE_ITEMS == front ) {
        front = (front + 1) % MAX_FTRACE_ITEMS;
    }
    
    FrameEntry* frame = &frame_table[rear];
    frame->src_addr = src_addr;
    frame->dst_addr = dst_addr;
    frame->src_func = find_func(src_addr);
    frame->dst_func = find_func(dst_addr);
    frame->type = type;
    frame->depth = depth;

    rear = (rear + 1) % MAX_FTRACE_ITEMS;
}

static void append_frame_call(vaddr_t src_addr, vaddr_t dst_addr) {
    append_frame(src_addr, dst_addr, FrameCall, depth++);
}

static void append_frame_return(vaddr_t src_addr, vaddr_t dst_addr) {
    append_frame(src_addr, dst_addr, FrameReturn, --depth);
}

static char* null_str = "<null?>";
static void pretty_print_frame(char* buf, FrameEntry* p) {
    static char* Action[] = { "Call", "Return" };
    char* src_func_name = p->src_func ? p->src_func->name : null_str;
    char* dst_func_name = p->dst_func ? p->dst_func->name : null_str;
    assert(p->type == FrameCall || p->type == FrameReturn);
    int n = 0;
    if (p->src_func == NULL)
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_RED), src_func_name, (paddr_t)p->src_addr);
    else if (p->src_func->level == LevelLow) 
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_BLUE), src_func_name, (paddr_t)p->src_addr);
    else if (p->src_func->level == LevelHigh) 
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_GREEN), src_func_name, (paddr_t)p->src_addr);

    n += sprintf(buf+n, "-(%s)->", Action[p->type]);

    if (p->dst_func == NULL)
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_RED), dst_func_name, (paddr_t)p->dst_addr);
    else if (p->dst_func->level == LevelLow) 
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_BLUE), dst_func_name, (paddr_t)p->dst_addr);
    else if (p->dst_func->level == LevelHigh) 
        n += sprintf(buf+n, ANSI_FMT("[%s@" FMT_PADDR "]", ANSI_FG_GREEN), dst_func_name, (paddr_t)p->dst_addr);        
    
    // sprintf(buf,
    //     "[%s@" FMT_PADDR "]" 
    //     "-(%s)->"
    //     "[%s@" FMT_PADDR "]", 
    //     src_func_name, (paddr_t)p->src_addr,
    //     Action[p->type], 
    //     dst_func_name, (paddr_t)p->dst_addr
    // );
}

/* 打印 func_table */
static void print_func_table() {
    printf(ANSI_FMT("===================================== <Function Table> =====================================\n", ANSI_FG_YELLOW));
    printf("%5s\t\t%16s\t\t%10s\t\t%5s\n", "No.", "name", "start", "size");
    for (int i = 0; i < func_count; ++i) {
        printf("%5d\t\t%16s\t\t%#10lx\t\t%5ld\t\n", i, func_table[i].name, func_table[i].start, func_table[i].size);
    }
    printf(ANSI_FMT("===================================== <Function Table> =====================================\n", ANSI_FG_YELLOW));
}

/* 打印 ftrace */
void ftrace_print() {
    
    print_func_table();

    char buf[1024];

    printf(ANSI_FMT("======================== <Function Trace> ========================\n", ANSI_FG_GREEN));

    size_t min_depth = SIZE_MAX;
    for (size_t i = 0; i < MAX_FTRACE_ITEMS; ++i) {
        if (frame_table[i].depth < min_depth) {
            min_depth = frame_table[i].depth;
        }
    }

    int i = front;
    while (i != rear) {
        int d = frame_table[i].depth - min_depth;
        switch (frame_table[i].type) {
            case FrameCall:
                pretty_print_frame(buf, &frame_table[i]);
                printf("%s%s\n", spaces(d*2), buf);
                depth += 1;
                break;
            case FrameReturn:
                depth -= 1;
                pretty_print_frame(buf, &frame_table[i]);
                printf("%s%s\n", spaces(d*2), buf);
                break;
            default: assert(0);
        }
        i = (i + 1) % MAX_FTRACE_ITEMS;
    }

    printf(ANSI_FMT("======================== <Function Trace> ========================\n", ANSI_FG_GREEN));
}

/* 初始化 ftrace */
void init_ftrace(const char* elf_file, const char* ramdisk_file, const char* appname) {
    if (elf_file == NULL) {
        printf("The elf file is not found.\n");
        return;
    }
    FILE* elf_fp = fopen(elf_file, "r");
    if (elf_fp == NULL) {
        printf("Can't open elf file %s.\n", elf_file);
        return;
    }
    
    fseek(elf_fp, 0, SEEK_END);
    size_t file_size = ftell(elf_fp);
    func_count = parse_elf(elf_fp, 0, file_size, func_table, LevelLow);

#ifdef CONFIG_LOAD_RAMDISK
    // /* copy from nanos-lite fs.c */
    static RamDiskFileInfo file_table[] = {
        #include "ramdisk.h" 
    };

    if (!(ramdisk_file && appname)) return;

	FILE* ramdisk_fp = fopen(ramdisk_file, "r");
	if (ramdisk_fp == NULL) {
		printf("Can't open ramdisk file %s.\n", ramdisk_file);
		return;
	}
    
	for (int i = 0; i < sizeof(file_table) / sizeof(file_table[0]); ++i) {
		if (strcmp(file_table[i].name, appname) != 0) continue;
		func_count += parse_elf(ramdisk_fp, file_table[i].disk_offset, file_table[i].size, func_table + func_count, LevelHigh);
		Log("The ramdisk is load, appname is %s", appname);
		return;
	}
#endif
}

void ftrace_inst_jalr(Decode* s) {
    uint32_t i = s->isa.inst.val;
    int rd  = BITS(i, 11, 7);
    int rs1 = BITS(i, 19, 15);
    int offset = SEXT(BITS(i, 31, 20), 12);
    // RETURN   ret  -> [jalr x0, 0(x1)]
    if (rd == 0 && rs1 == 1 && offset == 0) 
        append_frame_return(s->pc, s->dnpc);
    // CALL     call -> [jalr x1, imm(rs1)]
    else if (rd == 1)
        append_frame_call(s->pc, s->dnpc);
}

void ftrace_inst_jal(Decode* s) {
    uint32_t i = s->isa.inst.val;
    int rd  = BITS(i, 11, 7);
    // CALL     call -> [jal x1, offset]
    if (rd == 1) {
        append_frame_call(s->pc, s->dnpc);
    }
}
